import os
import matplotlib.pyplot as plt


def read_data_from_txt(file_path):
    y_coords = []
    confidences = []
    with open(file_path, "r") as file:
        for line in file:
            parts = line.strip().split()
            if len(parts) >= 6:
                y_coord = float(parts[3])
                confidence = float(parts[5])
                y_coords.append(y_coord)
                confidences.append(confidence)
    return y_coords, confidences


def collect_data_from_folder(folder_path):
    indexi = 0
    min_i = 0
    max_i = 100000
    all_y_coords = []
    all_confidences = []
    for file_name in os.listdir(folder_path):
        if file_name.endswith(".txt"):
            indexi = indexi + 1
            if indexi > min_i and indexi < max_i:
                file_path = os.path.join(folder_path, file_name)
                y_coords, confidences = read_data_from_txt(file_path)
                all_y_coords.extend(y_coords)
                all_confidences.extend(confidences)
    return all_y_coords, all_confidences


# 文件夹路径
model1_folder = "test_enhanced/labels"
model2_folder = "test_og/labels"

# 收集数据
model1_y_coords, model1_confidences = collect_data_from_folder(model1_folder)
model2_y_coords, model2_confidences = collect_data_from_folder(model2_folder)

# 绘制图像
plt.figure(figsize=(10, 6))
plt.scatter(model1_y_coords, model1_confidences, label="Model 1", alpha=0.5, marker="o")
plt.scatter(model2_y_coords, model2_confidences, label="Model 2", alpha=0.5, marker="x")
plt.xlabel("Y Coordinate")
plt.ylabel("Confidence")
plt.legend()
plt.show()
